import sys
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad

def encrypt_rbf(plain_rbf_path, encrypted_rbf_path, aes_key):
    """
    加密FPGA1的.rbf位流文件，生成V001.rbf
    :param plain_rbf_path: 原始.rbf文件路径
    :param encrypted_rbf_path: 加密后V001.rbf路径
    :param aes_key: AES-256密钥（32字节）
    """
    # 读取原始位流
    with open(plain_rbf_path, 'rb') as f:
        plain_data = f.read()
    
    # AES-256 CBC模式（需IV，这里用固定IV，实际可动态生成）
    iv = b'0123456789abcdef'  # 16字节IV
    cipher = AES.new(aes_key, AES.MODE_CBC, iv)
    
    # 填充数据（AES要求分组长度16字节）
    padded_data = pad(plain_data, AES.block_size)
    
    # 加密
    encrypted_data = cipher.encrypt(padded_data)
    
    # 写入加密文件（前16字节存IV，后续存密文）
    with open(encrypted_rbf_path, 'wb') as f:
        f.write(iv)
        f.write(encrypted_data)
    
    print(f"位流加密完成！输出文件：{encrypted_rbf_path}")
    print(f"加密后大小：{len(encrypted_data) + 16} 字节")

if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("用法：python rbf_encrypt.py <原始.rbf路径> <V001.rbf路径> <32字节AES密钥>")
        sys.exit(1)
    
    plain_rbf = sys.argv[1]
    encrypted_rbf = sys.argv[2]
    aes_key = sys.argv[3].encode('utf-8')
    
    if len(aes_key) != 32:
        print("错误：AES-256密钥必须为32字节！")
        sys.exit(1)
    
    encrypt_rbf(plain_rbf, encrypted_rbf, aes_key)